import json
import os
from functools import partial
from multiprocessing.pool import Pool
from pathlib import Path

import cv2
import pandas as pd
from tqdm import tqdm

from adaptive_contrast_enhancement.local_ace import adaptContrastEnhancement


def rgb_to_lace(rgb_path, lace_root, config, mode='train'):
    if mode == 'train':
        video_path = os.path.join(lace_root, 'train_release_LACE', Path(rgb_path).parent.name)
    # elif mode == 'dev':
    #     video_path = os.path.join(lace_root, 'Dev_files_LACE', Path(rgb_path).parent.name)
    else:
        video_path = os.path.join(lace_root, 'test_release_LACE', Path(rgb_path).parent.name)

    if not os.path.exists(video_path):
        try:
            os.mkdir(video_path)
        except OSError:
            pass
    lace_path = os.path.join(video_path, Path(rgb_path).name)

    if os.path.exists(os.path.join(video_path, Path(rgb_path).name)):
        return

    # 遍历该目录下的所有图片文件
    # print("filename :", rgb_path)
    img = cv2.imread(rgb_path)

    if img is None:
        print("The file name error,please check it")
        return -1

    img_lace = adaptContrastEnhancement(img, 15, 10)


    # cv2.imwrite('D:\\Datasets\\test_2'+ "\\"+filename,img_MSRCR)
    cv2.imwrite(lace_path, img_lace)


def get_images_paths(rgb_root, mode, data_csv):
    paths = []
    df = pd.read_csv(data_csv)
    for row in df.iterrows():
        if mode == 'train':
            path = os.path.join(rgb_root, 'train_release_crops', row[1]['video'], row[1]['file'])
        # elif mode == 'dev':
        #     path = os.path.join(rgb_root, 'Dev_files_crops', row[1]['video'], row[1]['file'])
        else:
            path = os.path.join(rgb_root, 'test_release_crops', row[1]['video'], row[1]['file'])
        paths.append(path)
    return paths


def main():
    with open('config_msr.json', 'r') as f:
        config = json.load(f)
    # mode = 'train'
    mode = 'test'

    # rgb_root = '/home/shaohua/data2/Datasets/Face_Anti_Spoofing/Oulu_NPU_5_frame'
    rgb_root = '/home/shaohua/data2/Datasets/Face_Anti_Spoofing/CASIA_FASD/casia'

    lace_root = '/home/shaohua/data2/Datasets/Face_Anti_Spoofing/CASIA_FASD/casia'
    # lace_root = '/home/shaohua/data2/Datasets/Face_Anti_Spoofing/Oulu_NPU'


    paths = get_images_paths(rgb_root, mode, '../data/casia_fasd/data_{}.csv'.format(mode))
    with Pool(processes=2) as p:
        with tqdm(total=len(paths)) as pbar:
            for v in p.imap_unordered(partial(rgb_to_lace, lace_root=lace_root, config=config, mode=mode), paths):
                pbar.update()



if __name__ == "__main__":
    main()
